# -*-coding: utf-8 -*-
"""
    @Author : panjq
    @E-mail : pan_jinquan@163.com
    @Date   : 2021-09-03 15:30:50
"""

import cv2
import numpy as np
import random
from detection.utils import image_utils, file_utils


class RandomImageMatting(object):
    """随机图像融合，需要提供RGB+Matte"""

    def __init__(self, p=1.0, bg_dir="bg_image/", is_rgb=False):
        """
        :param p: 概率
        :param bg_dir: 背景图库，PS：背景图不能含有目标检测的对象，避免背景的干扰
        """
        self.p = p
        self.bg_image_list = file_utils.get_files_lists(bg_dir, subname="")
        self.bg_nums = len(self.bg_image_list)
        self.is_rgb = is_rgb

    def random_read_bg_image(self, is_rgb=False, crop_rate=0.5):
        index = int(np.random.uniform(0, self.bg_nums))
        image_path = self.bg_image_list[index]
        # image_path = self.bg_image_list[0]
        image = cv2.imread(image_path)
        if is_rgb:
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        h_img, w_img, _ = image.shape
        xmin, ymin, xmax, ymax = self.extend_bboxes([0, 0, w_img, h_img], [0.8, 0.8])
        crop_xmin = int(random.uniform(0, xmin * crop_rate))
        crop_ymin = int(random.uniform(0, ymin * crop_rate))
        crop_xmax = int(min(w_img, random.uniform(w_img - xmax * crop_rate, w_img)))
        crop_ymax = int(min(h_img, random.uniform(h_img - ymax * crop_rate, h_img)))
        if random.random() < 0.5:
            image = image[:, ::-1, :]
        image = image[crop_ymin: crop_ymax, crop_xmin: crop_xmax]
        return image

    @staticmethod
    def extend_bboxes(box, scale=[1.0, 1.0]):
        """
        :param box: [xmin, ymin, xmax, ymax]
        :param scale: [sx,sy]==>(W,H)
        :return:
        """
        sx = scale[0]
        sy = scale[1]
        xmin, ymin, xmax, ymax = box[:4]
        cx = (xmin + xmax) / 2
        cy = (ymin + ymax) / 2

        ex_w = (xmax - xmin) * sx
        ex_h = (ymax - ymin) * sy
        ex_xmin = cx - 0.5 * ex_w
        ex_ymin = cy - 0.5 * ex_h
        ex_xmax = ex_xmin + ex_w
        ex_ymax = ex_ymin + ex_h
        ex_box = [ex_xmin, ex_ymin, ex_xmax, ex_ymax]
        return ex_box

    @staticmethod
    def image_alpha_blending(image, matte, bg_img=(219, 142, 67)):
        """
        图像融合
        更有效的C++实现: https://www.aiuai.cn/aifarm1237.html
        """
        h, w, d = image.shape
        if isinstance(bg_img, tuple) or isinstance(bg_img, list):
            bg = np.zeros_like(image, dtype=np.uint8)
            bg_img = np.asarray(bg[:, :, 0:3] + bg_img, dtype=np.uint8)
        if len(matte.shape) == 2:
            # alpha = cv2.cvtColor(alpha, cv2.COLOR_GRAY2BGR)
            matte = matte[:, :, np.newaxis]
        bg_img = cv2.resize(bg_img, dsize=(w, h))
        alpha = np.asarray(matte / 255.0, dtype=np.float32)
        image = image * alpha + bg_img * (1 - alpha)
        image = np.asarray(np.clip(image, 0, 255), dtype=np.uint8)
        return image

    def __call__(self, image, matte):
        """
        Args:
            image (numpy Image): Image to be cropped.
        Returns:
            PIL Image: Cropped image.
        """
        if random.random() < self.p:
            bg_image = self.random_read_bg_image(is_rgb=self.is_rgb)
            image = self.image_alpha_blending(image, matte, bg_img=bg_image)
        return image


if __name__ == "__main__":
    image_dir = "/home/dm/data3/dataset/segmentation/matting/1803151818/matting_00000000"
    bg_dir = "/home/dm/project/python-learning-notes/dataset/Test_Voc/JPEGImages"
    image_list = file_utils.get_files_lists(image_dir)
    matting = RandomImageMatting(bg_dir=bg_dir)
    for image_path in image_list:
        rgba = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
        rgb = rgba[:, :, 0:3]
        a = rgba[:, :, 3]
        dst = matting(rgb, a)
        image_utils.show_images_list("RGBA", [rgb, a, dst])
